Skip to content

feat(autojac): Use jac_to_grad to aggregate .jac fields #510

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
ValerianRey merged 37 commits into from
Jan 14, 2026
Merged

feat(autojac): Use jac_to_grad to aggregate .jac fields #510

ValerianRey merged 37 commits into from
Jan 14, 2026

Conversation

@ValerianRey
Copy link
Contributor

@ValerianRey ValerianRey commented Jan 8, 2026 •
edited
Loading

  • Move accumulation code to _accumulation.py
  • Rename Accumulate to AccumulateGrad
  • Add AccumulateJac
  • Add jac_to_grad
  • Remove aggregation responsibility from backward and mtl_backward
  • Remove Aggregate transform (move its code to jac_to_grad.py)
  • Simplify _disunite_gradient (use split instead of manually looping)
  • Add utils/asserts.py file to check things about .jac or .grad fields
  • Update tests and usage examples according to all those changes
  • Add changelog entry

@ValerianRey ValerianRey added the feat New feature or request label Jan 8, 2026
@ValerianRey ValerianRey self-assigned this Jan 8, 2026
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
@ValerianRey
Copy link
Contributor Author

ValerianRey commented Jan 8, 2026 •
edited
Loading

We have something a bit tricky: both jac_to_grad and the AccumulateGrad transform depend on the code from _accumulation.py.

To reuse some code in two different packages (utils and autojac), without making one depend on the other, we would need another protected package just for this code. So IMO we could:

  • Have a the _accumulation.py code somewhere else in torchjd (either in a standalone protected python file or in a protected accumulation module).
  • Duplicate some accumulation code.
  • Live with the fact that autojac depends on utils and utils depends on aggregation.

@ValerianRey
Copy link
Contributor Author

ValerianRey commented Jan 8, 2026

Lastly, it's very annoying to have this Pylance error all over the place:

Cannot access attribute "jac" for class "Tensor"
  Attribute "jac" is unknown

I'm not sure how we could reliably stop showing this error, or even use a custom tensor that would properly fix this.

@PierreQuinton
Copy link
Contributor

PierreQuinton commented Jan 8, 2026

Lastly, it's very annoying to have this Pylance error all over the place:

Cannot access attribute "jac" for class "Tensor"
  Attribute "jac" is unknown

I'm not sure how we could reliably stop showing this error, or even use a custom tensor that would properly fix this.

I think this is just because all this is bad design. We made it more torch compatible, and this is the price. Pylance is right, there is not field jac in Tensors.

@ValerianRey
Copy link
Contributor Author

ValerianRey commented Jan 8, 2026

I think this is just because all this is bad design. We made it more torch compatible, and this is the price. Pylance is right, there is not field jac in Tensors.

Yeah but there are still ways to make it much cleaner. I'm working on something.

@PierreQuinton
Copy link
Contributor

PierreQuinton commented Jan 8, 2026

Yeah but there are still ways to make it much cleaner. I'm working on something.

I think we should keep in mind that we would like the users to have a good typing experience. It would be nasty to make them use .jac fields if they get typing errors.

@ValerianRey
Copy link
Contributor Author

ValerianRey commented Jan 8, 2026

Yeah but there are still ways to make it much cleaner. I'm working on something.

I think we should keep in mind that we would like the users to have a good typing experience. It would be nasty to make them use .jac fields if they get typing errors.

Well, they should rarely need to use these .jac fields themselves. If they want access to the jacobians, they should use autojac.jac.

@PierreQuinton
Copy link
Contributor

PierreQuinton commented Jan 8, 2026

right, that's very good, and if they do, they should get some amount of warning. You can delete the discussion if you want (annoying that we cannot have a thread).

PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
PierreQuinton
PierreQuinton reviewed Jan 8, 2026
View reviewed changes
@ValerianRey ValerianRey marked this pull request as ready for review January 12, 2026 18:09
@ValerianRey
Remove duplicated optimizer.zero_grad() lines
e93f2d9
@ValerianRey
Fix formulation about freeing jacs
2fb6856
@ValerianRey
Add doc entry for jac_to_grad and usage example
57fe5b4
@ValerianRey
Add comments in jac_to_grad example
0a1fc21
@ValerianRey
Fix docstring of test_backward.py
f1ee074
@ValerianRey
Fix formatting in backward docstring
8b3d447
@ValerianRey
Fix comment in accumulate_jacs that applied to accumulate_grads
87b66f8
@ValerianRey
Fix error message in _check_expects_grad
1394395
@ValerianRey
Fix wrong import in basic_usage.rst
16349a0
@ValerianRey
Add explanation about how jac_to_grad works in jac_to_grad's docstring
0a8cc62
@ValerianRey
Improve description of parameters in jac_to_grad
c13a75b
@ValerianRey
Improve error message and usage example of jac_to_grad
674f6ad
@ValerianRey
Make _disunite_gradient use less memory
8a0fb0e
@ValerianRey
Free .jacs earlier to divide by two peak memory
0e8add2 
@PierreQuinton this was a big issue that we didn't spot earlier. I don't think the jacobian_matrix can be a view of the concatenated jacobians, so I think that having both the individual matrices + the combined matrix alive at the same time means using double memory. With this _free_jacs call much ealier, if the garbage collector is reactive, we shouldn't have this issue of doubling the peak memory usage for no reason. I think we should check that this PR doesn't introduce a huge memory efficiency regression. Can't merge without doing that.
ValerianRey
ValerianRey commented Jan 14, 2026
View reviewed changes
ValerianRey
ValerianRey commented Jan 14, 2026
View reviewed changes
@ValerianRey
Use Tensor.split in _disunit_gradient
430a8a2
@ValerianRey
Add kwargs to assert_jac_close and assert_grad_close
f0fe529 
These functions really are wrappers around assert_close, so we'd like them to always also take the parameters of assert_close, even if those change in the future, and to have the same default values. So I think kwargs is justified here. Also it's not user facing so the lack of documentation of the expected types will not be visible.
@ValerianRey
Rename expected_jacobian to J in some test
cff6d8e
@ValerianRey
Move asserts to tests/utils and use them in doc tests
84bd552
@ValerianRey
Rename test and update docstring to match its changes
4bb561d
@ValerianRey ValerianRey changed the title feat: Use .jac fields feat(autojac): Use jac_to_grad to aggregate .jac fields Jan 14, 2026
@ValerianRey ValerianRey merged commit 7e3637e into main Jan 14, 2026
13 checks passed
@ValerianRey ValerianRey deleted the revamp-interface branch January 14, 2026 15:43
@ValerianRey ValerianRey mentioned this pull request Jan 14, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@PierreQuinton PierreQuinton Awaiting requested review from PierreQuinton

Assignees

@ValerianRey ValerianRey

Labels

feat New feature or request package: autojac

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ValerianRey @PierreQuinton